package com.example.parttimejob.config;

import org.java_websocket.WebSocket;
import org.java_websocket.handshake.ClientHandshake;
import org.java_websocket.server.WebSocketServer;

import java.net.InetSocketAddress;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;

public class ChatWebSocketServer extends WebSocketServer {

    private final Map<WebSocket, String> userMap = new ConcurrentHashMap<>();
    private final Map<String, WebSocket> userConnections = new ConcurrentHashMap<>();
    private final Map<String, BlockingQueue<String>> offlineMessages = new ConcurrentHashMap<>();

    public ChatWebSocketServer(int port) {
        super(new InetSocketAddress(port));
    }

    @Override
    public void onOpen(WebSocket webSocket, ClientHandshake clientHandshake) {
        String username = userMap.get(webSocket);
        if (username != null) {
            userConnections.put(username, webSocket);
            webSocket.send("你的用户名为: " + username);
        }
    }

    @Override
    public void onClose(WebSocket conn, int code, String reason, boolean remote) {
        String username = userMap.get(conn);
        if (username != null) {
            userConnections.remove(username);
            broadcastMessage(username + " 离开了房间");
        }
    }

    @Override
    public void onMessage(WebSocket conn, String message) {
        String username = userMap.get(conn);

        if (username == null) {
            // First connection, treat the message as the username
            username = message;
            userMap.put(conn, username);
            if (offlineMessages.containsKey(username)) {
                BlockingQueue<String> messages = offlineMessages.get(username);
                while (!messages.isEmpty()) {
                    conn.send(messages.poll());
                }
                offlineMessages.remove(username);
            }
        } else {
            // Public message
            broadcastMessage(username + ": " + message);
        }
    }

    @Override
    public void onError(WebSocket conn, Exception ex) {
        // Handle the exception as needed
        ex.printStackTrace();
    }

    @Override
    public void onStart() {
        // Processing when the WebSocket server is started, if needed
    }

    private void broadcastMessage(String message) {
        Collection<WebSocket> connections = connections();
        for (WebSocket client : connections) {
            if (client.isOpen()) {
                client.send(message);
            }
        }
    }

    private Collection<WebSocket> connections() {
        return getConnections();
    }

    public static void main(String[] args) {
        // Use a dynamic port to avoid hardcoding the port number
        int port = 5678;
        if (args.length > 0) {
            try {
                port = Integer.parseInt(args[0]);
            } catch (NumberFormatException e) {
                System.err.println("没有端口: " + port);
            }
        }
        new ChatWebSocketServer(port).start();
    }
}